from libs.PipeLine import PipeLine, ScopedTiming
from libs.AIBase import AIBase
from libs.AI2D import Ai2d
import os
import ujson
from media.media import *
import nncase_runtime as nn
import ulab.numpy as np
import image
import gc
import sys
import aidemo
import time




class PersonKeyPointApp(AIBase):
    # 保持原有代码不变
    def __init__(self, kmodel_path, model_input_size, confidence_threshold=0.1, nms_threshold=0.5, rgb888p_size=[1280, 720], display_size=[1920, 1080], debug_mode=0):
        super().__init__(kmodel_path, model_input_size, rgb888p_size, debug_mode)
        self.kmodel_path = kmodel_path
        self.model_input_size = model_input_size
        self.confidence_threshold = confidence_threshold
        self.nms_threshold = nms_threshold
        self.rgb888p_size = [ALIGN_UP(rgb888p_size[0], 16), rgb888p_size[1]]
        self.display_size = [ALIGN_UP(display_size[0], 16), display_size[1]]
        self.debug_mode = debug_mode
        self.SKELETON = [(16, 14), (14, 12), (17, 15), (15, 13), (12, 13), (6, 12), (7, 13), (6, 7), (6, 8), (7, 9), (8, 10), (9, 11), (2, 3), (1, 2), (1, 3), (2, 4), (3, 5), (4, 6), (5, 7)]
        self.LIMB_COLORS = [(255, 51, 153, 255), (255, 51, 153, 255), (255, 51, 153, 255), (255, 51, 153, 255), (255, 255, 51, 255), (255, 255, 51, 255), (255, 255, 51, 255), (255, 255, 128, 0), (255, 255, 128, 0), (255, 255, 128, 0), (255, 255, 128, 0), (255, 255, 128, 0), (255, 0, 255, 0), (255, 0, 255, 0), (255, 0, 255, 0), (255, 0, 255, 0), (255, 0, 255, 0), (255, 0, 255, 0), (255, 0, 255, 0)]
        self.KPS_COLORS = [(255, 0, 255, 0), (255, 0, 255, 0), (255, 0, 255, 0), (255, 0, 255, 0), (255, 0, 255, 0), (255, 255, 128, 0), (255, 255, 128, 0), (255, 255, 128, 0), (255, 255, 128, 0), (255, 255, 128, 0), (255, 255, 128, 0), (255, 51, 153, 255), (255, 51, 153, 255), (255, 51, 153, 255), (255, 51, 153, 255), (255, 51, 153, 255), (255, 51, 153, 255)]
        self.ai2d = Ai2d(debug_mode)
        self.ai2d.set_ai2d_dtype(nn.ai2d_format.NCHW_FMT, nn.ai2d_format.NCHW_FMT, np.uint8, np.uint8)

    def config_preprocess(self, input_image_size=None):
        with ScopedTiming("set preprocess config", self.debug_mode > 0):
            ai2d_input_size = input_image_size if input_image_size else self.rgb888p_size
            top, bottom, left, right = self.get_padding_param()
            self.ai2d.pad([0, 0, 0, 0, top, bottom, left, right], 0, [0, 0, 0])
            self.ai2d.resize(nn.interp_method.tf_bilinear, nn.interp_mode.half_pixel)
            self.ai2d.build([1, 3, ai2d_input_size[1], ai2d_input_size[0]], [1, 3, self.model_input_size[1], self.model_input_size[0]])

    def postprocess(self, results):
        with ScopedTiming("postprocess", self.debug_mode > 0):
            results = aidemo.person_kp_postprocess(results[0], [self.rgb888p_size[1], self.rgb888p_size[0]], self.model_input_size, self.confidence_threshold, self.nms_threshold)
            return results

    def draw_result(self, pl, res):
        with ScopedTiming("display_draw", self.debug_mode > 0):
            if res[0]:
                pl.osd_img.clear()
                kpses = res[1]
                for i in range(len(res[0])):
                    for k in range(17 + 2):
                        if k < 17:
                            kps_x, kps_y, kps_s = round(kpses[i][k][0]), round(kpses[i][k][1]), kpses[i][k][2]
                            kps_x1 = int(float(kps_x) * self.display_size[0] // self.rgb888p_size[0])
                            kps_y1 = int(float(kps_y) * self.display_size[1] // self.rgb888p_size[1])
                            if kps_s > 0:
                                pl.osd_img.draw_circle(kps_x1, kps_y1, 5, self.KPS_COLORS[k], 4)
                        if k < len(self.SKELETON):
                            ske = self.SKELETON[k]
                            pos1_x, pos1_y = round(kpses[i][ske[0]-1][0]), round(kpses[i][ske[0]-1][1])
                            pos1_x_ = int(float(pos1_x) * self.display_size[0] // self.rgb888p_size[0])
                            pos1_y_ = int(float(pos1_y) * self.display_size[1] // self.rgb888p_size[1])

                            pos2_x, pos2_y = round(kpses[i][(ske[1]-1)][0]), round(kpses[i][(ske[1]-1)][1])
                            pos2_x_ = int(float(pos2_x) * self.display_size[0] // self.rgb888p_size[0])
                            pos2_y_ = int(float(pos2_y) * self.display_size[1] // self.rgb888p_size[1])

                            pos1_s, pos2_s = kpses[i][(ske[0]-1)][2], kpses[i][(ske[1]-1)][2]
                            if pos1_s > 0.0 and pos2_s > 0.0:
                                pl.osd_img.draw_line(pos1_x_, pos1_y_, pos2_x_, pos2_y_, self.LIMB_COLORS[k], 4)
                    pl.osd_img.draw_circle(kps_x1, kps_y1, 3, (0, 255, 0), 2)  # 绿色圆心表示肩臀中点
                    gc.collect()
            else:
                pl.osd_img.clear()

    def get_padding_param(self):
        dst_w = self.model_input_size[0]
        dst_h = self.model_input_size[1]
        input_width = self.rgb888p_size[0]
        input_high = self.rgb888p_size[1]
        ratio_w = dst_w / input_width
        ratio_h = dst_h / input_high
        ratio = min(ratio_w, ratio_h)
        new_w = int(ratio * input_width)
        new_h = int(ratio * input_high)
        dw = (dst_w - new_w) // 2
        dh = (dst_h - new_h) // 2
        return dh, dh, dw, dw
